Using pre-trained OCP models as ASE calculators
Contents
Using pre-trained OCP models as ASE calculators#
For those interested in using our pretrained models for other applications, we provide an [ASE](https://wiki.fysik.dtu.dk/ase/#:~:text=The%20Atomic%20Simulation%20Environment%20(ASE,under%20the%20GNU%20LGPL%20license.)-compatible Calculator to interface with ASE’s functionality.
Download pretrained checkpoint#
We have released checkpoints of all the models on the leaderboard here. These trained models can be used as an ASE calculator for various calculations.
For this tutorial we download our current best model checkpoint: GemNet-T
!wget -q -nc https://dl.fbaipublicfiles.com/opencatalystproject/models/2021_08/s2ef/gemnet_t_direct_h512_all.pt
checkpoint_path = "gemnet_t_direct_h512_all.pt"
Using the OCP Calculator#
import os
import ase.io
import numpy as np
from ase.build import add_adsorbate, fcc100, molecule
from ase.constraints import FixAtoms
from ase.optimize import BFGS
from ocpmodels.common.relaxation.ase_utils import OCPCalculator
# Construct a sample structure
adslab = fcc100("Cu", size=(3, 3, 3))
adsorbate = molecule("C3H8")
add_adsorbate(adslab, adsorbate, 3, offset=(1, 1))
tags = np.zeros(len(adslab))
tags[18:27] = 1
tags[27:] = 2
adslab.set_tags(tags)
cons = FixAtoms(indices=[atom.index for atom in adslab if (atom.tag == 0)])
adslab.set_constraint(cons)
adslab.center(vacuum=13.0, axis=2)
adslab.set_pbc(True)
# Define the calculator
calc = OCPCalculator(checkpoint=checkpoint_path)
# Set up the calculator
adslab.calc = calc
os.makedirs("data/sample_ml_relax", exist_ok=True)
opt = BFGS(adslab, trajectory="data/sample_ml_relax/toy_c3h8_relax.traj")
opt.run(fmax=0.05, steps=100)
amp: false
cmd:
checkpoint_dir: /home/runner/work/ml_catalysis_tutorials/ml_catalysis_tutorials/ml-catalysis-tutorials/notes/checkpoints/2022-10-30-17-23-12
commit: cba9fb6
identifier: ''
logs_dir: /home/runner/work/ml_catalysis_tutorials/ml_catalysis_tutorials/ml-catalysis-tutorials/notes/logs/tensorboard/2022-10-30-17-23-12
print_every: 100
results_dir: /home/runner/work/ml_catalysis_tutorials/ml_catalysis_tutorials/ml-catalysis-tutorials/notes/results/2022-10-30-17-23-12
seed: null
timestamp_id: 2022-10-30-17-23-12
dataset: null
gpus: 0
logger: tensorboard
model: gemnet_t
model_attributes:
activation: silu
cbf:
name: spherical_harmonics
cutoff: 6.0
direct_forces: true
emb_size_atom: 512
emb_size_bil_trip: 64
emb_size_cbf: 16
emb_size_edge: 512
emb_size_rbf: 16
emb_size_trip: 64
envelope:
exponent: 5
name: polynomial
extensive: true
max_neighbors: 50
num_after_skip: 2
num_atom: 3
num_before_skip: 1
num_blocks: 3
num_concat: 1
num_radial: 128
num_spherical: 7
otf_graph: true
output_init: HeOrthogonal
rbf:
name: gaussian
regress_forces: true
noddp: false
optim:
batch_size: 32
clip_grad_norm: 10
ema_decay: 0.999
energy_coefficient: 1
eval_batch_size: 32
eval_every: 5000
factor: 0.8
force_coefficient: 100
loss_energy: mae
loss_force: l2mae
lr_initial: 0.0005
max_epochs: 80
mode: min
num_workers: 2
optimizer: AdamW
optimizer_params:
amsgrad: true
patience: 3
scheduler: ReduceLROnPlateau
slurm: {}
task:
dataset: trajectory_lmdb
description: Regressing to energies and forces for DFT trajectories from OCP
eval_on_free_atoms: true
grad_input: atomic forces
labels:
- potential energy
metric: mae
train_on_free_atoms: true
type: regression
trainer: forces
/home/runner/work/ml_catalysis_tutorials/ml_catalysis_tutorials/ocp/ocpmodels/preprocessing/atoms_to_graphs.py:139: UserWarning: Creating a tensor from a list of numpy.ndarrays is extremely slow. Please consider converting the list to a single numpy.ndarray with numpy.array() before converting to a tensor. (Triggered internally at /opt/conda/conda-bld/pytorch_1639180507909/work/torch/csrc/utils/tensor_new.cpp:201.)
cell = torch.Tensor(atoms.get_cell()).view(1, 3, 3)
/home/runner/micromamba-root/envs/buildenv/lib/python3.9/site-packages/torch/autocast_mode.py:141: UserWarning: User provided device_type of 'cuda', but CUDA is not available. Disabling
warnings.warn('User provided device_type of \'cuda\', but CUDA is not available. Disabling')
/home/runner/micromamba-root/envs/buildenv/lib/python3.9/site-packages/torch/functional.py:1069: UserWarning: torch.meshgrid: in an upcoming release, it will be required to pass the indexing argument. (Triggered internally at /opt/conda/conda-bld/pytorch_1639180507909/work/aten/src/ATen/native/TensorShape.cpp:2157.)
return _VF.cartesian_prod(tensors) # type: ignore[attr-defined]
/home/runner/work/ml_catalysis_tutorials/ml_catalysis_tutorials/ocp/ocpmodels/models/gemnet/gemnet.py:373: UserWarning: __floordiv__ is deprecated, and its behavior will change in a future version of pytorch. It currently rounds toward 0 (like the 'trunc' function NOT 'floor'). This results in incorrect rounding for negative values. To keep the current behavior, use torch.div(a, b, rounding_mode='trunc'), or for actual floor division, use torch.div(a, b, rounding_mode='floor').
neighbors_new // 2,
/home/runner/work/ml_catalysis_tutorials/ml_catalysis_tutorials/ocp/ocpmodels/models/gemnet/gemnet.py:467: UserWarning: __floordiv__ is deprecated, and its behavior will change in a future version of pytorch. It currently rounds toward 0 (like the 'trunc' function NOT 'floor'). This results in incorrect rounding for negative values. To keep the current behavior, use torch.div(a, b, rounding_mode='trunc'), or for actual floor division, use torch.div(a, b, rounding_mode='floor').
block_sizes = neighbors // 2
Step Time Energy fmax
BFGS: 0 17:23:40 -4.099784 1.5675
BFGS: 1 17:23:40 -4.244461 1.1370
BFGS: 2 17:23:41 -4.403120 0.7635
BFGS: 3 17:23:42 -4.503653 0.8364
BFGS: 4 17:23:43 -4.558208 0.7339
BFGS: 5 17:23:44 -4.592068 0.4095
BFGS: 6 17:23:44 -4.619359 0.7312
BFGS: 7 17:23:45 -4.671457 0.9712
BFGS: 8 17:23:46 -4.796439 0.9212
BFGS: 9 17:23:47 -4.957980 0.9762
BFGS: 10 17:23:47 -5.109450 1.0385
BFGS: 11 17:23:48 -5.295599 1.2249
BFGS: 12 17:23:49 -5.498955 1.1272
BFGS: 13 17:23:50 -5.618106 1.0669
BFGS: 14 17:23:50 -5.737125 0.9508
BFGS: 15 17:23:51 -5.901956 0.9261
BFGS: 16 17:23:52 -6.076142 1.2739
BFGS: 17 17:23:52 -6.198369 1.2028
BFGS: 18 17:23:53 -6.250305 0.6852
BFGS: 19 17:23:54 -6.254100 0.2008
BFGS: 20 17:23:55 -6.293970 0.1779
BFGS: 21 17:23:55 -6.326323 0.2295
BFGS: 22 17:23:56 -6.324457 0.1700
BFGS: 23 17:23:57 -6.321297 0.1016
BFGS: 24 17:23:57 -6.328427 0.0847
BFGS: 25 17:23:58 -6.331800 0.0587
BFGS: 26 17:23:59 -6.332060 0.0444
True
import matplotlib.pyplot as plt
from ase.visualize.plot import animate
from matplotlib import rc
# the `index` argument corresponds to what frame of the trajectory to read in, specifiying ":" reads in the full trajectory.
traj = ase.io.read("data/sample_ml_relax/toy_c3h8_relax.traj", index=":")
anim = animate(traj, radii=0.8, rotation=("-75x, 45y, 10z"))
rc("animation", html="jshtml")
anim
import pandas as pd
import plotly.express as px
energies = [image.get_potential_energy() for image in traj]
df = pd.DataFrame(
{
"Relaxation Step": range(len(energies)),
"GemNet-T Adsorption Energy [eV]": energies,
}
)
fig = px.line(df, x="Relaxation Step", y="GemNet-T Adsorption Energy [eV]")
fig.show()